# vim:ts=4:sw=4:noet:
""" Intermediate Compiler Representation """

import spy
import cfg
from environment import Environment
import cmg
import logging
from c.base.support import CoreException

# IR Node
class IRNode:
	def __init__(self, line):
		self.line = line
		self.checked = False

	def cfg(self, body, bb, inBlock=False):
		raise CoreException("Unhandled CFG: %s" % self)

	def cg(self, cgen):
		raise CoreException("Unhandled CG: %s" % self)

	def genConstLit(self, cgen):
		raise CoreException("Unhandled genConstLit: %s" % self)

	def isJump(self):
		return False
	
	def isSuper(self):
		return False

# IR Class
class IRClass(IRNode):
	def __init__(self, line, name, supername, ivars, meths, docs, pragmas):
		IRNode.__init__(self, line)
		self.name = name
		self.supername = supername
		self.ivars = ivars
		self.meths = meths
		self.docs = docs
		self.pragmas = pragmas

# IR Method
class IRMethod(IRNode):
	def __init__(self, line, name, args, temps, maxtemps, body, source, docs):
		IRNode.__init__(self, line)
		self.name = name
		self.args = args
		self.temps = temps
		self.maxtemps = maxtemps
		self.body = body
		self.source = source
		self.docs = docs

	def getSig(self):
		if len(self.args)==1:
			return self.name
		elif len(self.args)==2 and ":" not in self.name:
			return "%s %s" % (self.name, self.args[1])
		else:
			r = ""
			for i, k in enumerate(self.name.split(":")):
				if not k:
					break
				r += "%s: %s " % (k, self.args[i+1])
			return r[:-1]

	def canInline(self):
		cannotInline = [
				"ifTrue:", 
				"ifFalse:", 
				"ifTrue:ifFalse:",
				"whileTrue:",
				"whileFalse:",
				]
		if ((self.name in cannotInline) or
				(len(self.body.statements)!=1) or
				(not isinstance(self.body.statements[0], IRReturn))):
			return False
		retv = self.body.statements[0].value
		return isinstance(retv, (IRInteger,IRString,IRChar,IRSymbol))

	def cfg(self):
		self.body.cfg()
		
	def cg(self, mem, cgen):
		cgen.lineNum(self.line)
		self.body.cg(cgen)

# IR Body
class IRBody(IRNode):
	def __init__(self, line, statements):
		IRNode.__init__(self, line)
		self.statements = statements
		self.bb = None

	def cfg(self):
		self.bb = cfg.CFGBasicBlock(name="entry")
		curbb = self.bb
		for st in self.statements:
			curbb = st.cfg(self, curbb)
			curbb.addStatement(cfg.CFGEndStatement())
		curbb.addStatement(cfg.CFGEndBody())
			
	def cg(self, cgen):
		cgen.lineNum(self.line)
		self.bb.cleanup()
		self.bb.cg(cgen)

# IR Block
class IRBlock(IRNode):
	def __init__(self, line, args, statements):
		IRNode.__init__(self, line)
		self.statements = statements
		self.args = args
		self.bb = None
	
	def cfg(self, body, bb, inline=False, inBlock=False):
		if inline:
			curbb = bb
			for i, st in enumerate(self.statements):
				curbb = st.cfg(self, curbb, inBlock)
				if i<(len(self.statements)-1):
					curbb.addStatement(cfg.CFGEndStatement())
			return curbb
		else:
			curbb = cfg.CFGBasicBlock(name="block")
			startbb = curbb
			bb.addNext(curbb)
			for i, st in enumerate(self.statements):
				curbb = st.cfg(self, curbb, True)
				if i<(len(self.statements)-1):
					curbb.addStatement(cfg.CFGEndStatement())
			if len(self.statements)==0:
				curbb.addStatement(IRName(self.line, "nil", 0))
			curbb.addStatement(cfg.CFGEndBlock())
			nextbb = cfg.CFGBasicBlock(name="endblock")
			bb.addStatement(cfg.CFGStartBlock(self.tempStart, startbb, nextbb))
			curbb.addNext(nextbb)
			bb.addNext(nextbb)
			return nextbb
	
	def cg(self, cgen):
		pass

	def __repr__(self):
		return "(irblock %d %d)" % (len(self.statements), len(self.args))

# IR Message
class IRMessage(IRNode):
	def __init__(self, line, recv, name, args):
		IRNode.__init__(self, line)
		self.name = name
		self.recv = recv
		self.args = args

	def cfgIf(self, body, bb, truen, falsen, inBlock):
		if self.recv:
			bb = self.recv.cfg(body, bb, inBlock)
		bbt = cfg.CFGBasicBlock(name=truen)
		bbf = cfg.CFGBasicBlock(name=falsen)
		bbn = cfg.CFGBasicBlock(name="endif")
		bb.addNext(bbt)
		bb.addNext(bbf)
		return bb, bbt, bbf, bbn

	def cfgSpecialIf(self, body, bb, truen, falsen, nextTrue, nextVal, inBlock):
		bb, bbt, bbf, bbn = self.cfgIf(body, bb, truen, falsen, inBlock)
		if nextTrue:
			bba, bbb = bbt, bbf
		else:
			bba, bbb = bbf, bbt
		bbb.addNext(bbn)
		bbb.addStatement(IRName(self.line, nextVal, 0))
		self.args[0].cfg(body, bba, inline=True, inBlock=inBlock).addNext(bbn)
		if len(self.args[0].statements)==0:
			bba.addStatement(IRName(self.line, "nil", 0))
		bb.addStatement(cfg.CFGIf(bbt, bbf, bbn, iftrue=(not nextTrue)))
		return bbn

	def cfg(self, body, bb, inBlock=False):
		argsAreBlocks = all(isinstance(arg, IRBlock) for arg in self.args)
		if self.name=="ifTrue:ifFalse:" and argsAreBlocks:
			# optimise a ifTrue: [...] ifFalse: [...] calls
			bb, bbt, bbf, bbn = self.cfgIf(body, bb, "true", "false", inBlock)
			self.args[0].cfg(body, bbt, inline=True, inBlock=inBlock).addNext(bbn)
			self.args[1].cfg(body, bbf, inline=True, inBlock=inBlock).addNext(bbn)
			if len(self.args[0].statements)==0:
				bbt.addStatement(IRName(self.line, "nil", 0))
			if len(self.args[1].statements)==0:
				bbf.addStatement(IRName(self.line, "nil", 0))
			bb.addStatement(cfg.CFGIf(bbt, bbf, bbn))
			return bbn
		elif self.name=="ifTrue:" and argsAreBlocks:
			# optimise a ifTrue: [...] call
			return self.cfgSpecialIf(body, bb, "true", "false", True, "nil", inBlock)
		elif self.name=="ifFalse:" and argsAreBlocks:
			# optimise a ifFalse: [...] call
			return self.cfgSpecialIf(body, bb, "true", "false", False, "nil", inBlock)
		elif self.name=="and:" and argsAreBlocks:
			# optimise a and: [...] call
			return self.cfgSpecialIf(body, bb, "and", "shortcirc", True, "false", inBlock)
		elif self.name=="or:" and argsAreBlocks:
			# optimise a or: [...] call
			return self.cfgSpecialIf(body, bb, "shortcirc", "or", False, "true", inBlock)
		elif self.name=="whileTrue:" and isinstance(self.recv, IRBlock) and argsAreBlocks:
			# optimise a whileTrue: [...] call
			bbc = cfg.CFGBasicBlock(name="whilecond")
			bbw = cfg.CFGBasicBlock(name="whilebody")
			bbn = cfg.CFGBasicBlock(name="endwhile")
			bb.addNext(bbc)
			bbcend = self.recv.cfg(body, bbc, inline=True)
			bbcend.addNext(bbw)
			bbcend.addNext(bbn)
			bbwend = self.args[0].cfg(body, bbw, inline=True, inBlock=inBlock)
			if len(self.args[0].statements)>0:
				bbwend.addStatement(cfg.CFGEndStatement())
			bbwend.addNext(bbc)
			bbn.addStatement(IRName(self.line, "nil", 0))
			bbcend.addStatement(cfg.CFGIf(bbw, bbn, bbn, iftrue=False))
			return bbn
		else:
			if self.recv:
				bb = self.recv.cfg(body, bb, inBlock=inBlock)
			for arg in self.args:
				bb = arg.cfg(body, bb, inBlock=inBlock)
			bb.addStatement(self)
		return bb
	
	def cg(self, cgen):
		cgen.lineNum(self.line)
		unames = ["isNil", "notNil"]
		symname = cgen.mem.newSymbol(self.name)
		if self.recv and self.recv.isSuper():
			cgen.sendSuper(symname, len(self.args))
		elif len(self.args)==1 and self.name in cmg.binmsgs:
			cgen.sendb(symname, cmg.binmsgs.index(self.name))
		elif len(self.args)==0 and self.name in unames:
			cgen.sendu(symname, unames.index(self.name))
		else:
			cgen.send(symname, len(self.args))
	
	def __repr__(self):
		return "(irmsg %s %d)" % (self.name, len(self.args))

# IR Return
class IRReturn(IRNode):
	def __init__(self, line, value):
		IRNode.__init__(self, line)
		self.value = value

	def cfg(self, body, bb, inBlock=False):
		self.inBlock = inBlock
		bb = self.value.cfg(body, bb, inBlock)
		bb.addStatement(self)
		return bb

	def cg(self, cgen):
		cgen.lineNum(self.line)
		if self.inBlock:
			cgen.returnBlock()
		else:
			cgen.returnStack()

	def __repr__(self):
		return "(irreturn)"

# IR Integer
class IRInteger(IRNode):
	def __init__(self, line, value):
		IRNode.__init__(self, line)
		self.value = value
	
	def cfg(self, body, bb, inBlock=False):
		bb.addStatement(self)
		return bb

	def cg(self, cgen):
		cgen.lineNum(self.line)
		if (self.value < 8192) and (self.value >= -8192):
			cgen.pushLiteral(spy.SpySmallInt(cgen.mem, self.value))
		else:
			cgen.pushLiteral(spy.SpyInteger(cgen.mem, self.value))

	def genConstLit(self, cgen):
		if (self.value < 8192) and (self.value >= -8192):
			return spy.SpySmallInt(cgen.mem, self.value)
		else:
			return spy.SpyInteger(cgen.mem, self.value)

	def __repr__(self):
		return "(irint %d)" % (self.value)

# IR String
class IRString(IRNode):
	def __init__(self, line, value):
		IRNode.__init__(self, line)
		self.value = value

	def cfg(self, body, bb, inBlock=False):
		bb.addStatement(self)
		return bb

	def cg(self, cgen):
		cgen.lineNum(self.line)
		cgen.pushLiteral(spy.SpyString(cgen.mem, self.value))

	def genConstLit(self, cgen):
		return spy.SpyString(cgen.mem, self.value)

	def __repr__(self):
		return "(irstr '%s')" % (self.value)

# IR Char
class IRChar(IRNode):
	def __init__(self, line, value):
		IRNode.__init__(self, line)
		self.value = value

	def cfg(self, body, bb, inBlock=False):
		bb.addStatement(self)
		return bb

	def cg(self, cgen):
		cgen.lineNum(self.line)
		cgen.pushLiteral(spy.SpyChar(cgen.mem, self.value))

	def genConstLit(self, cgen):
		return spy.SpyChar(cgen.mem, self.value)

	def __repr__(self):
		return "(irchar '%s')" % (self.value)

# IR Name
class IRName(IRNode):
	def __init__(self, line, name, type=None):
		IRNode.__init__(self, line)
		self.name = name
		if type!=None:
			self.type = type

	def cfg(self, body, bb, inBlock=False):
		bb.addStatement(self)
		return bb

	def cg(self, cgen):
		cgen.lineNum(self.line)
		if self.type==Environment.GLOBAL:
			if self.name=="Smalltalk":
				cgen.pushSmalltalk()
			else:
				if not cgen.mem.getGlobal(self.name):
					logging.error("couldn't find global '%s'" % self.name)
				cgen.pushLiteral(cgen.mem.getGlobal(self.name))
		elif self.type==Environment.IVAR:
			cgen.pushInstance(self.pos)
		elif self.type==Environment.METH_TEMP:
			cgen.pushTemporary(self.pos)
		elif self.type==Environment.BLOCK_ARG:
			cgen.pushTemporary(self.pos)
		else:
			raise CoreException("Unhandled irname type in cg: %s" % Environment.typenames[self.type])
	
	def isSuper(self):
		return self.type==Environment.METH_TEMP and self.pos==-1
	
	def __repr__(self):
		return "(irname '%s')" % self.name

# IR Symbol
class IRSymbol(IRNode):
	def __init__(self, line, name):
		IRNode.__init__(self, line)
		self.name = name

	def cfg(self, body, bb, inBlock=False):
		bb.addStatement(self)
		return bb
	
	def cg(self, cgen):
		cgen.lineNum(self.line)
		cgen.pushLiteral(cgen.mem.newSymbol(self.name))
	
	def genConstLit(self, cgen):
		return spy.SpySymbol(cgen.mem, self.name)

	def __repr__(self):
		return "(irsym #%s)" % (self.name)

# IR Array (Const)
class IRArray(IRNode):
	def __init__(self, line, values):
		IRNode.__init__(self, line)
		self.values = values

	def cfg(self, body, bb, inBlock=False):
		bb.addStatement(self)
		return bb
	
	def cg(self, cgen):
		cgen.lineNum(self.line)
		cgen.pushLiteral(
				spy.SpyArray(cgen.mem, 
					map(lambda x: x.genConstLit(cgen), self.values)))

	def __repr__(self):
		return "(irarray #(%s))" % (self.values)

# IR Primitive
class IRPrimitive(IRNode):
	def __init__(self, line, value, args):
		IRNode.__init__(self, line)
		self.value = value
		self.args = args

	def cfg(self, body, bb, inBlock=False):
		for arg in self.args:
			bb = arg.cfg(body, bb, inBlock)
		bb.addStatement(self)
		return bb
	
	def cg(self, cgen):
		cgen.lineNum(self.line)
		cgen.prim(self.value, len(self.args))
	
	def __repr__(self):
		return "(irprim %d %d)" % (self.value, len(self.args))

# IR Assign
class IRAssign(IRNode):
	def __init__(self, line, var, value):
		IRNode.__init__(self, line)
		self.var = var
		self.value = value
		self.doPop = False

	def addPop(self):
		self.doPop = True
	
	def cfg(self, body, bb, inBlock=False):
		bb = self.value.cfg(body, bb, inBlock)
		bb.addStatement(self)
		return bb
	
	def cg(self, cgen):
		cgen.lineNum(self.line)
		if self.var.type==Environment.IVAR:
			if self.doPop:
				cgen.popInstance(self.var.pos)
			else:
				cgen.storeInstance(self.var.pos)
		elif (self.var.type==Environment.METH_TEMP or
					self.var.type==Environment.BLOCK_ARG):
			if self.doPop:
				cgen.popTemporary(self.var.pos)
			else:
				cgen.storeTemporary(self.var.pos)
		elif (self.var.type==Environment.GLOBAL and 
					self.var.name=="Smalltalk"):
			cgen.popSmalltalk()
		else:
			raise CoreException("Unhandled irname type in cg: %s (%s)" % (
				Environment.typenames[self.var.type], self.var.name))
	
	def __repr__(self):
		return "(irassign %s)" % (self.var)

# IR Cascade
class IRCascade(IRNode):
	def __init__(self, line, node):
		IRNode.__init__(self, line)
		self.node  = node

	def cfg(self, body, bb, inBlock=False):
		self.node.cfg(body, bb, inBlock)
		for i, casc in enumerate(self.list):
			last = i == (len(self.list)-1)
			if not last:
				bb.addStatement(IRCascadeMark(1))
				casc.cfg(body, bb, inBlock)
				bb.addStatement(IRCascadeMark(0))
			else:
				casc.cfg(body, bb, inBlock)
		return bb

class IRCascadeMark:
	def __init__(self, start):
		self.start = start

	def cg(self, cgen):
		if self.start:
			cgen.cascadeStart()
		else:
			cgen.cascadeEnd()

	def __repr__(self):
		return "(ircascade %d)" % self.start
